#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import random

sys.path.append('.')

"""
   对FaceID进行合并
"""


def main():
    parser = argparse.ArgumentParser(description="FaceID Merge")
    parser.add_argument(
        "--faceid_file_list", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/FaceID/typical_file_faceid_list.txt", help="path to faceid file list", type=str
    )
    parser.add_argument(
        "--faceid_merge_file_list", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/FaceID/select_image_label_merge.txt", help="path to merge faceid file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/FaceID", help="path to output", type=str
    )
    args = parser.parse_args()

    random.seed(123456)

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    merge_faceid_videos = {}
    with open(args.faceid_merge_file_list, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4be6600/video/20190926_161748/00013_000.jpg 2789
            lines = line.strip().split()
            file_name = lines[0]
            video_name = os.path.dirname(file_name)
            faceid = int(lines[1])
            video_names = merge_faceid_videos.get(faceid, None)
            if video_names is None:
                merge_faceid_videos[faceid] = [video_name]
            else:
                video_names.append(video_name)

    file_name_faceid_list = []
    video_faceids = {}
    with open(args.faceid_file_list, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4be6600/video/20190926_161748/00013_000.jpg 2789
            lines = line.strip().split()
            file_name = lines[0]
            video_name = os.path.dirname(file_name)
            faceid = int(lines[1])
            file_name_faceid_list.append([file_name, faceid])
            video_faceids[video_name] = faceid

    faceid_map = {}
    for merge_faceid, video_names in merge_faceid_videos.items():
        for video_name in video_names:
            faceid = video_faceids[video_name]
            faceid_map[faceid] = merge_faceid

    out_file_path = os.path.join(output_dir, 'M'+os.path.basename(args.faceid_file_list))
    with open(out_file_path, 'w') as file:
        file_name_faceid_list.sort(key=lambda x: x[0])
        for file_name, faceid in file_name_faceid_list:
            faceid = faceid_map[faceid]
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    print('FaceID Count:', len(merge_faceid_videos.keys()))
    print('finish!')


if __name__ == '__main__':
    main()
